#include "bw_tree.h"
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <assert.h>

// =============================================================================
// Node Traversal and Search Utilities
// =============================================================================

// Find position in leaf node where key should be inserted
static int find_leaf_position(bw_node_t* node, const void* key, size_t key_size, bool* found) {
    if (!node || node->type != BW_NODE_LEAF || !key || !found) {
        if (found) *found = false;
        return -1;
    }
    
    *found = false;
    char* data_ptr = node->data;
    int left = 0, right = node->count - 1;
    
    // Binary search for key position
    while (left <= right) {
        int mid = (left + right) / 2;
        
        // Calculate offset to mid-th key-value pair
        size_t offset = 0;
        for (int i = 0; i < mid; i++) {
            bw_kv_pair_t* pair = (bw_kv_pair_t*)(data_ptr + offset);
            offset += sizeof(bw_kv_pair_t) + pair->key_size + pair->value_size;
        }
        
        bw_kv_pair_t* pair = (bw_kv_pair_t*)(data_ptr + offset);
        int cmp = bw_key_compare_raw(key, key_size, pair->data, pair->key_size);
        
        if (cmp == 0) {
            *found = true;
            return mid;
        } else if (cmp < 0) {
            right = mid - 1;
        } else {
            left = mid + 1;
        }
    }
    
    return left; // Position where key should be inserted
}

// Find child page ID in internal node
static page_id_t find_child_page(bw_node_t* node, const void* key, size_t key_size) {
    if (!node || node->type != BW_NODE_INTERNAL || !key) {
        return 0;
    }
    
    char* data_ptr = node->data;
    size_t offset = 0;
    
    // Internal node format: [key1, page_id1, key2, page_id2, ..., page_id_n]
    // First page_id is for keys < key1, last is for keys >= last_key
    
    for (int i = 0; i < node->count; i++) {
        // Read key
        uint16_t key_len = *(uint16_t*)(data_ptr + offset);
        offset += sizeof(uint16_t);
        
        char* node_key = data_ptr + offset;
        offset += key_len;
        
        // Read page ID
        page_id_t page_id = *(page_id_t*)(data_ptr + offset);
        offset += sizeof(page_id_t);
        
        int cmp = bw_key_compare_raw(key, key_size, node_key, key_len);
        if (cmp < 0) {
            return page_id;
        }
    }
    
    // Key is >= all keys, return last page ID
    page_id_t last_page_id = *(page_id_t*)(data_ptr + offset);
    return last_page_id;
}

// Apply all delta records to get the logical view of a node
static bw_node_t* get_logical_node_view(bw_node_t* node_chain, memory_manager_t* memory_mgr) {
    if (!node_chain || !memory_mgr) return NULL;
    
    // Find base node by traversing delta chain
    bw_node_t* base_node = node_chain;
    int delta_count = 0;
    
    while (base_node && base_node->type >= BW_NODE_DELTA_INSERT) {
        base_node = atomic_load(&base_node->next);
        delta_count++;
        
        // Prevent infinite loops
        if (delta_count > BW_TREE_MAX_DELTA_CHAIN * 2) {
            return NULL;
        }
    }
    
    if (!base_node || (base_node->type != BW_NODE_LEAF && base_node->type != BW_NODE_INTERNAL)) {
        return NULL;
    }
    
    // If no deltas, return base node
    if (delta_count == 0) {
        return base_node;
    }
    
    // Create consolidated node by applying deltas
    // This is a simplified version - full implementation would handle all delta types
    bw_node_t* logical_node = (base_node->type == BW_NODE_LEAF) ?
        bw_node_create_leaf(memory_mgr, base_node->size) :
        bw_node_create_internal(memory_mgr, base_node->size);
    
    if (!logical_node) return NULL;
    
    // Copy base node data
    memcpy(logical_node->data, base_node->data, 
           base_node->size - sizeof(bw_node_t));
    logical_node->count = base_node->count;
    logical_node->level = base_node->level;
    
    // Apply deltas in reverse order (oldest to newest)
    // This is a simplified implementation for demonstration
    
    return logical_node;
}

// Check if node needs consolidation
bool bw_node_needs_consolidation(bw_node_t* node, uint32_t max_delta_chain) {
    if (!node) return false;
    
    int delta_count = 0;
    bw_node_t* current = node;
    
    while (current && current->type >= BW_NODE_DELTA_INSERT) {
        delta_count++;
        if (delta_count >= max_delta_chain) {
            return true;
        }
        current = atomic_load(&current->next);
    }
    
    return false;
}

// =============================================================================
// Core Operations Implementation
// =============================================================================

bw_result_t bw_tree_search(bw_tree_t* tree, const void* key, size_t key_size,
                          bw_search_result_t* result) {
    if (!tree || !key || key_size == 0 || !result || atomic_load(&tree->is_destroyed)) {
        return BW_ERROR_INVALID_PARAM;
    }
    
    epoch_t epoch = gc_enter_epoch(tree->gc_context);
    atomic_fetch_add(&tree->active_operations, 1);
    atomic_fetch_add(&tree->total_operations, 1);
    
    result->found = false;
    result->value = NULL;
    result->page_id = 0;
    result->node = NULL;
    result->index = 0;
    
    page_id_t current_page_id = tree->root_page_id;
    
    // Traverse tree from root to leaf
    while (current_page_id != 0) {
        bw_page_t* page = mapping_table_get_page(tree->mapping_table, current_page_id);
        if (!page) {
            atomic_fetch_sub(&tree->active_operations, 1);
            gc_exit_epoch(tree->gc_context);
            return BW_ERROR_INVALID_PARAM;
        }
        
        bw_node_t* node = atomic_load(&page->node_ptr);
        if (!node) {
            atomic_fetch_sub(&page->ref_count, 1);
            atomic_fetch_sub(&tree->active_operations, 1);
            gc_exit_epoch(tree->gc_context);
            return BW_ERROR_INVALID_PARAM;
        }
        
        // Get logical view of node (apply deltas)
        bw_node_t* logical_node = get_logical_node_view(node, tree->memory_mgr);
        if (!logical_node) {
            atomic_fetch_sub(&page->ref_count, 1);
            atomic_fetch_sub(&tree->active_operations, 1);
            gc_exit_epoch(tree->gc_context);
            return BW_ERROR_INVALID_PARAM;
        }
        
        if (logical_node->type == BW_NODE_LEAF) {
            // Search in leaf node
            bool found = false;
            int position = find_leaf_position(logical_node, key, key_size, &found);
            
            if (found) {
                // Extract value from leaf node
                char* data_ptr = logical_node->data;
                size_t offset = 0;
                
                for (int i = 0; i < position; i++) {
                    bw_kv_pair_t* pair = (bw_kv_pair_t*)(data_ptr + offset);
                    offset += sizeof(bw_kv_pair_t) + pair->key_size + pair->value_size;
                }
                
                bw_kv_pair_t* pair = (bw_kv_pair_t*)(data_ptr + offset);
                
                // Create value copy
                result->value = bw_value_create(tree->memory_mgr, 
                                              pair->data + pair->key_size, 
                                              pair->value_size);
                result->found = true;
                result->page_id = current_page_id;
                result->node = logical_node;
                result->index = position;
            }
            
            if (logical_node != node) {
                bw_node_destroy(logical_node, tree->memory_mgr);
            }
            
            atomic_fetch_sub(&page->ref_count, 1);
            break;
        } else if (logical_node->type == BW_NODE_INTERNAL) {
            // Find child page in internal node
            page_id_t child_page_id = find_child_page(logical_node, key, key_size);
            
            if (logical_node != node) {
                bw_node_destroy(logical_node, tree->memory_mgr);
            }
            
            atomic_fetch_sub(&page->ref_count, 1);
            current_page_id = child_page_id;
        } else {
            // Invalid node type
            if (logical_node != node) {
                bw_node_destroy(logical_node, tree->memory_mgr);
            }
            atomic_fetch_sub(&page->ref_count, 1);
            atomic_fetch_sub(&tree->active_operations, 1);
            gc_exit_epoch(tree->gc_context);
            return BW_ERROR_INVALID_PARAM;
        }
    }
    
    atomic_fetch_sub(&tree->active_operations, 1);
    gc_exit_epoch(tree->gc_context);
    
    return result->found ? BW_SUCCESS : BW_ERROR_KEY_NOT_FOUND;
}

bw_result_t bw_tree_insert(bw_tree_t* tree, const void* key, size_t key_size,
                          const void* value, size_t value_size) {
    if (!tree || !key || key_size == 0 || !value || value_size == 0 ||
        key_size > BW_TREE_MAX_KEY_SIZE || value_size > BW_TREE_MAX_VALUE_SIZE ||
        atomic_load(&tree->is_destroyed)) {
        return BW_ERROR_INVALID_PARAM;
    }
    
    epoch_t epoch = gc_enter_epoch(tree->gc_context);
    atomic_fetch_add(&tree->active_operations, 1);
    atomic_fetch_add(&tree->total_operations, 1);
    
    bw_result_t result = BW_SUCCESS;
    page_id_t target_page_id = 0;
    
    // Find target leaf page
    page_id_t current_page_id = tree->root_page_id;
    
    while (current_page_id != 0) {
        bw_page_t* page = mapping_table_get_page(tree->mapping_table, current_page_id);
        if (!page) {
            result = BW_ERROR_INVALID_PARAM;
            goto cleanup;
        }
        
        bw_node_t* node = atomic_load(&page->node_ptr);
        if (!node) {
            atomic_fetch_sub(&page->ref_count, 1);
            result = BW_ERROR_INVALID_PARAM;
            goto cleanup;
        }
        
        bw_node_t* logical_node = get_logical_node_view(node, tree->memory_mgr);
        if (!logical_node) {
            atomic_fetch_sub(&page->ref_count, 1);
            result = BW_ERROR_INVALID_PARAM;
            goto cleanup;
        }
        
        if (logical_node->type == BW_NODE_LEAF) {
            target_page_id = current_page_id;
            
            if (logical_node != node) {
                bw_node_destroy(logical_node, tree->memory_mgr);
            }
            atomic_fetch_sub(&page->ref_count, 1);
            break;
        } else if (logical_node->type == BW_NODE_INTERNAL) {
            page_id_t child_page_id = find_child_page(logical_node, key, key_size);
            
            if (logical_node != node) {
                bw_node_destroy(logical_node, tree->memory_mgr);
            }
            atomic_fetch_sub(&page->ref_count, 1);
            current_page_id = child_page_id;
        } else {
            if (logical_node != node) {
                bw_node_destroy(logical_node, tree->memory_mgr);
            }
            atomic_fetch_sub(&page->ref_count, 1);
            result = BW_ERROR_INVALID_PARAM;
            goto cleanup;
        }
    }
    
    if (target_page_id == 0) {
        result = BW_ERROR_INVALID_PARAM;
        goto cleanup;
    }
    
    // Attempt to insert into target leaf page using delta record
    bw_page_t* target_page = mapping_table_get_page(tree->mapping_table, target_page_id);
    if (!target_page) {
        result = BW_ERROR_INVALID_PARAM;
        goto cleanup;
    }
    
    // Create insert delta record
    delta_record_t* insert_delta = bw_delta_create_insert(tree->memory_mgr, 
                                                         key, key_size, 
                                                         value, value_size);
    if (!insert_delta) {
        atomic_fetch_sub(&target_page->ref_count, 1);
        result = BW_ERROR_OUT_OF_MEMORY;
        goto cleanup;
    }
    
    // Try to install delta record using CAS
    bw_node_t* current_head = atomic_load(&target_page->node_ptr);
    atomic_store(&insert_delta->next, current_head);
    
    if (!atomic_compare_exchange_strong(&target_page->node_ptr, &current_head, 
                                       (bw_node_t*)insert_delta)) {
        // CAS failed, another thread modified the node
        bw_node_destroy((bw_node_t*)insert_delta, tree->memory_mgr);
        atomic_fetch_sub(&target_page->ref_count, 1);
        result = BW_ERROR_CONCURRENT_UPDATE;
        goto cleanup;
    }
    
    atomic_fetch_add(&tree->total_nodes, 1);
    
    // Check if consolidation is needed
    if (bw_node_needs_consolidation((bw_node_t*)insert_delta, tree->max_delta_chain)) {
        // Trigger consolidation (async)
        bw_node_consolidate(tree, target_page_id);
    }
    
    atomic_fetch_sub(&target_page->ref_count, 1);
    
cleanup:
    atomic_fetch_sub(&tree->active_operations, 1);
    gc_exit_epoch(tree->gc_context);
    return result;
}

bw_result_t bw_tree_delete(bw_tree_t* tree, const void* key, size_t key_size) {
    if (!tree || !key || key_size == 0 || atomic_load(&tree->is_destroyed)) {
        return BW_ERROR_INVALID_PARAM;
    }
    
    epoch_t epoch = gc_enter_epoch(tree->gc_context);
    atomic_fetch_add(&tree->active_operations, 1);
    atomic_fetch_add(&tree->total_operations, 1);
    
    bw_result_t result = BW_SUCCESS;
    page_id_t target_page_id = 0;
    
    // Find target leaf page (similar to insert)
    page_id_t current_page_id = tree->root_page_id;
    
    while (current_page_id != 0) {
        bw_page_t* page = mapping_table_get_page(tree->mapping_table, current_page_id);
        if (!page) {
            result = BW_ERROR_INVALID_PARAM;
            goto cleanup;
        }
        
        bw_node_t* node = atomic_load(&page->node_ptr);
        if (!node) {
            atomic_fetch_sub(&page->ref_count, 1);
            result = BW_ERROR_INVALID_PARAM;
            goto cleanup;
        }
        
        bw_node_t* logical_node = get_logical_node_view(node, tree->memory_mgr);
        if (!logical_node) {
            atomic_fetch_sub(&page->ref_count, 1);
            result = BW_ERROR_INVALID_PARAM;
            goto cleanup;
        }
        
        if (logical_node->type == BW_NODE_LEAF) {
            target_page_id = current_page_id;
            
            if (logical_node != node) {
                bw_node_destroy(logical_node, tree->memory_mgr);
            }
            atomic_fetch_sub(&page->ref_count, 1);
            break;
        } else if (logical_node->type == BW_NODE_INTERNAL) {
            page_id_t child_page_id = find_child_page(logical_node, key, key_size);
            
            if (logical_node != node) {
                bw_node_destroy(logical_node, tree->memory_mgr);
            }
            atomic_fetch_sub(&page->ref_count, 1);
            current_page_id = child_page_id;
        } else {
            if (logical_node != node) {
                bw_node_destroy(logical_node, tree->memory_mgr);
            }
            atomic_fetch_sub(&page->ref_count, 1);
            result = BW_ERROR_INVALID_PARAM;
            goto cleanup;
        }
    }
    
    if (target_page_id == 0) {
        result = BW_ERROR_KEY_NOT_FOUND;
        goto cleanup;
    }
    
    // Create delete delta record and install it
    bw_page_t* target_page = mapping_table_get_page(tree->mapping_table, target_page_id);
    if (!target_page) {
        result = BW_ERROR_INVALID_PARAM;
        goto cleanup;
    }
    
    delta_record_t* delete_delta = bw_delta_create_delete(tree->memory_mgr, key, key_size);
    if (!delete_delta) {
        atomic_fetch_sub(&target_page->ref_count, 1);
        result = BW_ERROR_OUT_OF_MEMORY;
        goto cleanup;
    }
    
    // Install delete delta using CAS
    bw_node_t* current_head = atomic_load(&target_page->node_ptr);
    atomic_store(&delete_delta->next, current_head);
    
    if (!atomic_compare_exchange_strong(&target_page->node_ptr, &current_head,
                                       (bw_node_t*)delete_delta)) {
        bw_node_destroy((bw_node_t*)delete_delta, tree->memory_mgr);
        atomic_fetch_sub(&target_page->ref_count, 1);
        result = BW_ERROR_CONCURRENT_UPDATE;
        goto cleanup;
    }
    
    atomic_fetch_add(&tree->total_nodes, 1);
    
    // Check if consolidation is needed
    if (bw_node_needs_consolidation((bw_node_t*)delete_delta, tree->max_delta_chain)) {
        bw_node_consolidate(tree, target_page_id);
    }
    
    atomic_fetch_sub(&target_page->ref_count, 1);
    
cleanup:
    atomic_fetch_sub(&tree->active_operations, 1);
    gc_exit_epoch(tree->gc_context);
    return result;
}

bw_result_t bw_tree_update(bw_tree_t* tree, const void* key, size_t key_size,
                          const void* value, size_t value_size) {
    if (!tree || !key || key_size == 0 || !value || value_size == 0 ||
        key_size > BW_TREE_MAX_KEY_SIZE || value_size > BW_TREE_MAX_VALUE_SIZE ||
        atomic_load(&tree->is_destroyed)) {
        return BW_ERROR_INVALID_PARAM;
    }
    
    // For simplicity, implement update as delete + insert
    // A full implementation would use update delta records
    
    bw_result_t delete_result = bw_tree_delete(tree, key, key_size);
    if (delete_result != BW_SUCCESS && delete_result != BW_ERROR_KEY_NOT_FOUND) {
        return delete_result;
    }
    
    return bw_tree_insert(tree, key, key_size, value, value_size);
}

// Conditional operations
bw_result_t bw_tree_insert_if_not_exists(bw_tree_t* tree, const void* key, size_t key_size,
                                        const void* value, size_t value_size) {
    if (!tree || !key || key_size == 0 || !value || value_size == 0 ||
        atomic_load(&tree->is_destroyed)) {
        return BW_ERROR_INVALID_PARAM;
    }
    
    // Check if key exists
    bw_search_result_t search_result;
    bw_result_t result = bw_tree_search(tree, key, key_size, &search_result);
    
    if (result == BW_SUCCESS) {
        // Key exists, clean up and return error
        if (search_result.value) {
            bw_value_destroy(search_result.value, tree->memory_mgr);
        }
        return BW_ERROR_KEY_EXISTS;
    } else if (result == BW_ERROR_KEY_NOT_FOUND) {
        // Key doesn't exist, proceed with insert
        return bw_tree_insert(tree, key, key_size, value, value_size);
    } else {
        // Other error
        return result;
    }
}

bw_result_t bw_tree_update_if_exists(bw_tree_t* tree, const void* key, size_t key_size,
                                    const void* value, size_t value_size) {
    if (!tree || !key || key_size == 0 || !value || value_size == 0 ||
        atomic_load(&tree->is_destroyed)) {
        return BW_ERROR_INVALID_PARAM;
    }
    
    // Check if key exists
    bw_search_result_t search_result;
    bw_result_t result = bw_tree_search(tree, key, key_size, &search_result);
    
    if (result == BW_SUCCESS) {
        // Key exists, clean up search result and proceed with update
        if (search_result.value) {
            bw_value_destroy(search_result.value, tree->memory_mgr);
        }
        return bw_tree_update(tree, key, key_size, value, value_size);
    } else {
        // Key doesn't exist or other error
        return result;
    }
}

bw_result_t bw_tree_delete_if_exists(bw_tree_t* tree, const void* key, size_t key_size) {
    if (!tree || !key || key_size == 0 || atomic_load(&tree->is_destroyed)) {
        return BW_ERROR_INVALID_PARAM;
    }
    
    // Attempt delete - if key doesn't exist, treat as success
    bw_result_t result = bw_tree_delete(tree, key, key_size);
    return (result == BW_ERROR_KEY_NOT_FOUND) ? BW_SUCCESS : result;
}